#include "zab_communication.h"
#include "arpa/inet.h"
#include <stdlib.h>
#include <assert.h>
#include <stddef.h>

void zab_packetizer_init(ZAB_PACKETIZER_T* packetizer, ZAB_PACKET_HANDLER_T handler, void* user_data)
{
    packetizer->packet = NULL;
    packetizer->packet_length = 0;
    packetizer->packet_type = 0;
    packetizer->user_data = user_data;
    packetizer->handler = handler;
    packetizer->state = PACKETIZE_GETTING_TYPE;
    packetizer->index = 0;
}

void zab_packetizer_push(ZAB_PACKETIZER_T* packetizer, unsigned char* data, int length)
{
    int i;
    for (i = 0; i < length; i++)
    {
        unsigned char c = data[i];
        if (packetizer->state == PACKETIZE_GETTING_TYPE)
        {
            packetizer->buffer[packetizer->index++] = c;
            if (packetizer->index == 2) // sizeof(unsigned short)
            {
                packetizer->packet_type = ntohs(*(unsigned short*)packetizer->buffer);
                packetizer->state = PACKETIZE_GETTING_LENGTH;
                packetizer->index = 0;
            }
        }
        else if (packetizer->state == PACKETIZE_GETTING_LENGTH)
        {
            packetizer->buffer[packetizer->index++] = c;
            if (packetizer->index == 4) // sizeof(unsigned int)
            {
                packetizer->packet_length = ntohl(*(unsigned int*)packetizer->buffer);
                packetizer->state = PACKETIZE_GETTING_DATA;
                packetizer->index = 0;
                // create the packet to hold data
                packetizer->packet = malloc(packetizer->packet_length);
                assert(packetizer->packet != NULL);
                // note: still in net byte order
                packetizer->packet->type = htons(packetizer->packet_type);
                packetizer->packet->length = htonl(packetizer->packet_length);
                // special case:
                // this packet does not hold any data, only header is available
                if (packetizer->packet_length == 6)
                {
                    // we have received a full packet, run the callback
                    packetizer->handler(packetizer->packet, packetizer->user_data);
                    // destroy the packet
                    free(packetizer->packet);
                    zab_packetizer_init(packetizer, packetizer->handler, packetizer->user_data);
                }
            }
        }
        else if (packetizer->state == PACKETIZE_GETTING_DATA)
        {
            packetizer->packet->data[packetizer->index++] = c;
            if (packetizer->index == packetizer->packet_length - 6)
            {
                // we have received a full packet, run the callback
                packetizer->handler(packetizer->packet, packetizer->user_data);
                // destroy the packet
                free(packetizer->packet);
                zab_packetizer_init(packetizer, packetizer->handler, packetizer->user_data);
            }
        }
    }
}

void zab_packetizer_push_byte(ZAB_PACKETIZER_T* packetizer, unsigned char byte)
{
    zab_packetizer_push(packetizer, &byte, 1);
}

ZAB_PACKET_T* zab_vote_msg_encode(ZAB_VOTE_MSG_T* msg, PACKET_BUFFER_T* pb)
{
    packet_buffer_clear(pb);
    packet_buffer_push_uint16(pb, ZAB_PACKET_VOTE);
    packet_buffer_push_uint32(pb, 6 + sizeof(*msg));
    packet_buffer_push_uint32(pb, msg->node_id);
    packet_buffer_push_uint32(pb, msg->propsed_leader);
    packet_buffer_push_uint32(pb, msg->proposed_zxid.epoch);
    packet_buffer_push_uint32(pb, msg->proposed_zxid.count);
    packet_buffer_push_uint32(pb, msg->election_round);
    packet_buffer_push_uint32(pb, msg->node_role);
    packet_buffer_push_uint32(pb, msg->node_epoch);

    assert(packet_buffer_get_uint32(pb, 2) == (unsigned int)pb->data_length);
    return (ZAB_PACKET_T*)pb->buffer;
}

void zab_vote_msg_decode(ZAB_VOTE_MSG_T* msg, unsigned char* data, int length)
{
    PACKET_BUFFER_T pb;
    unsigned short type;
    unsigned int len;

    packet_buffer_parse(&pb, data, length);
    packet_buffer_shift_uint16(&pb, &type);
    packet_buffer_shift_uint32(&pb, &len);
    packet_buffer_shift_uint32(&pb, &msg->node_id);
    packet_buffer_shift_uint32(&pb, &msg->propsed_leader);
    packet_buffer_shift_uint32(&pb, &msg->proposed_zxid.epoch);
    packet_buffer_shift_uint32(&pb, &msg->proposed_zxid.count);
    packet_buffer_shift_uint32(&pb, &msg->election_round);
    packet_buffer_shift_uint32(&pb, &msg->node_role);
    packet_buffer_shift_uint32(&pb, &msg->node_epoch);
}



ZAB_PACKET_T* zab_new_epoch_msg_encode(ZAB_NEW_EPOCH_MSG_T* msg, PACKET_BUFFER_T* pb)
{
    packet_buffer_clear(pb);
    packet_buffer_push_uint16(pb, ZAB_PACKET_NEW_EPOCH);
    packet_buffer_push_uint32(pb, 6 + sizeof(ZAB_NEW_EPOCH_MSG_T));
    packet_buffer_push_uint32(pb, msg->node_id);
    packet_buffer_push_uint32(pb, msg->epoch);
    packet_buffer_push_uint32(pb, msg->zxid.epoch);
    packet_buffer_push_uint32(pb, msg->zxid.count);
    return (ZAB_PACKET_T*)pb->buffer;
}

void zab_new_epoch_msg_decode(ZAB_NEW_EPOCH_MSG_T* msg, unsigned char* data, int length)
{
    PACKET_BUFFER_T pb;
    unsigned short type;
    unsigned int len;
    packet_buffer_parse(&pb, data, len);

    packet_buffer_shift_uint16(&pb, &type);
    packet_buffer_shift_uint16(&pb, &len);
    packet_buffer_shift_uint32(&pb, &msg->node_id);
    packet_buffer_shift_uint32(&pb, &msg->epoch);
    packet_buffer_shift_uint32(&pb, &msg->zxid.epoch);
    packet_buffer_shift_uint32(&pb, &msg->zxid.count);
}

ZAB_PACKET_T* zab_new_epoch_ack_msg_encode(ZAB_NEW_EPOCH_ACK_MSG_T* msg, PACKET_BUFFER_T* pb)
{
    packet_buffer_clear(pb);
    packet_buffer_push_uint16(pb, ZAB_PACKET_NEW_EPOCH);
    packet_buffer_push_uint32(pb, 6 + sizeof(ZAB_NEW_EPOCH_ACK_MSG_T));
    packet_buffer_push_uint32(pb, msg->node_id);
    packet_buffer_push_uint32(pb, msg->epoch);
    packet_buffer_push_uint32(pb, msg->zxid.epoch);
    packet_buffer_push_uint32(pb, msg->zxid.count);
    return (ZAB_PACKET_T*)pb->buffer;
}

void zab_new_epoch_ack_msg_decode(ZAB_NEW_EPOCH_ACK_MSG_T* msg, unsigned char* data, int length)
{
    PACKET_BUFFER_T pb;
    unsigned short type;
    unsigned int len;
    packet_buffer_parse(&pb, data, len);

    packet_buffer_shift_uint16(&pb, &type);
    packet_buffer_shift_uint16(&pb, &len);
    packet_buffer_shift_uint32(&pb, &msg->node_id);
    packet_buffer_shift_uint32(&pb, &msg->epoch);
    packet_buffer_shift_uint32(&pb, &msg->zxid.epoch);
    packet_buffer_shift_uint32(&pb, &msg->zxid.count);
}

ZAB_PACKET_T* zab_new_leader_msg_encode(ZAB_NEW_LEADER_MSG_T* msg, PACKET_BUFFER_T* pb)
{
    packet_buffer_clear(pb);
    packet_buffer_push_uint16(pb, ZAB_PACKET_NEW_EPOCH);
    packet_buffer_push_uint32(pb, 6 + sizeof(ZAB_NEW_LEADER_MSG_T));
    packet_buffer_push_uint32(pb, msg->leader_node);
    packet_buffer_push_uint32(pb, msg->zxid.epoch);
    packet_buffer_push_uint32(pb, msg->zxid.count);
    return (ZAB_PACKET_T*)pb->buffer;
}

void zab_new_leader_msg_decode(ZAB_NEW_LEADER_MSG_T* msg, unsigned char* data, int length)
{
    PACKET_BUFFER_T pb;
    unsigned short type;
    unsigned int len;
    packet_buffer_parse(&pb, data, len);

    packet_buffer_shift_uint16(&pb, &type);
    packet_buffer_shift_uint16(&pb, &len);
    packet_buffer_shift_uint32(&pb, &msg->leader_node);
    packet_buffer_shift_uint32(&pb, &msg->zxid.epoch);
    packet_buffer_shift_uint32(&pb, &msg->zxid.count);
}

ZAB_PACKET_T* zab_new_leader_ack_msg_encode(ZAB_NEW_LEADER_ACK_MSG_T* msg, PACKET_BUFFER_T* pb)
{
    packet_buffer_clear(pb);
    packet_buffer_push_uint16(pb, ZAB_PACKET_NEW_EPOCH);
    packet_buffer_push_uint32(pb, 6 + sizeof(ZAB_NEW_LEADER_MSG_T));
    packet_buffer_push_uint32(pb, msg->leader_node);
    packet_buffer_push_uint32(pb, msg->zxid.epoch);
    packet_buffer_push_uint32(pb, msg->zxid.count);
    return (ZAB_PACKET_T*)pb->buffer;
}

void zab_new_leader_ack_msg_decode(ZAB_NEW_LEADER_ACK_MSG_T* msg, unsigned char* data, int length)
{
    PACKET_BUFFER_T pb;
    unsigned short type;
    unsigned int len;
    packet_buffer_parse(&pb, data, len);

    packet_buffer_shift_uint16(&pb, &type);
    packet_buffer_shift_uint16(&pb, &len);
    packet_buffer_shift_uint32(&pb, &msg->leader_node);
    packet_buffer_shift_uint32(&pb, &msg->zxid.epoch);
    packet_buffer_shift_uint32(&pb, &msg->zxid.count);
}
